from langchain_core.messages import HumanMessage, AIMessage

from com.wp.langchain.exercise.LLMInit import LLMInit

llm = LLMInit().get_llm()
response = llm.invoke([HumanMessage(content="你好，我是wupeng，我来自中国")])
print(response)

invoke = llm.invoke([HumanMessage(content="我的名字叫什么")])
print(invoke)

llm_invoke = llm.invoke(
    [HumanMessage(content="Hi! I'm WuPeng"), AIMessage(content="Hello Bob! How can I assist you today?"),
     HumanMessage(content="What's my name?"), ])
print(llm_invoke)

from langchain_core.chat_history import (
    BaseChatMessageHistory,
    InMemoryChatMessageHistory,
)
from langchain_core.runnables.history import RunnableWithMessageHistory

store = {}


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory
    return store[session_id]


with_message_history = RunnableWithMessageHistory(llm, get_session_history)

config = {"configurable": {"session_id": "abc2"}}
history_invoke = with_message_history.invoke([HumanMessage(content="Hi! I'm WuPeng")], config=config)
print(history_invoke.content)

response = with_message_history.invoke([HumanMessage(content="What's my name?")], config=config)

print(response.content)
